#include "ch_task.hpp"

uint32_t check_grease(uint32_t num){
    if( (num & 0x0a0a) != 0x0a0a )
        return 0;
    if( ((num & 0xf000) >> 12) != (num & 0x000f) )
        return 0;
    return 1;
}

void run_ch_task(std::shared_ptr<ch_pkt> p, thread_safe_queue<ch_info> * queue){

    #ifdef DEBUG
        std::cout << "run_ch_task" << std::endl;
    #endif

    // len("255.255.255.255") = 3 * 4 + 3 = 15
    char client_ip[INET_ADDRSTRLEN];
    char server_ip[INET_ADDRSTRLEN];
    // handshake protocol version
    char ver_str[VERSION_LEN];
    char ciphers_str[CIPHERS_LEN];
    char exts_str[EXTS_LEN];
    // support groups
    char supgs_str[SUPPG_LEN];
    // support ec point format
    char suppfs_str[SUPPFS_LEN];
    char ticket_str[TICKET_LEN];
    char sni_str[SNI_LEN];
    char ja3_str[JA3_LEN];
    char ja3_md5[SMALL_STR_LEN];
    // 16进制的 ticket_str
    char ticket_0X_str[TICKET_LEN * 2];
    // unix时间
    char unix_time[SMALL_STR_LEN];

    uint32_t ver_str_len = 0;
    uint32_t ciphers_str_len = 0;
    uint32_t exts_str_len = 0;
    uint32_t supgs_str_len = 0;
    uint32_t suppfs_str_len = 0;
    uint32_t ticket_str_len = 0;
    uint32_t sni_str_len = 0;
    uint32_t ja3_str_len = 0;
    uint32_t unix_time_len = 0;

    // 读报文的 IP 部分
    struct iphdr * ip = (struct iphdr *)(p.get()->raw_data + ETHER_HDR_LEN);

    // 判断是否是IP报文，如果不是IP报文的话直接丢弃
    if( (ip->ihl) * 4 < 20 )
        return;

    // 判断是否发生了IP报文分片，发生了IP分片的话，不考虑重组IP报文，直接丢弃
    if( (UINT16_NTOP(ip->frag_off)) != 0x4000 )
        return;

    // 读报文的 TCP 部分
    struct tcphdr * tcp = (struct tcphdr *)(p.get()->raw_data + ETHER_HDR_LEN + sizeof(struct iphdr));

    // 判断是否是TCP报文，如果不是TCP报文的话直接丢弃
    if( (tcp->th_off) * 4 < 20)
        return;

    // 读报文的 TLS 部分中的具有固定格式的前半部分。先判断包长度是否足够，避免访问到非法内存区域
    struct ch_pkt_header * ch;
    if( (p.get()->pkt_info).caplen > ETHER_HDR_LEN + sizeof(struct iphdr) + sizeof(struct tcphdr) + sizeof(struct ch_pkt_header) ){
        ch = (struct ch_pkt_header *)(p.get()->raw_data + ETHER_HDR_LEN + sizeof(struct iphdr) + sizeof(struct tcphdr));
    }

    // 报文中已被读取的字节数
    uint32_t pkt_read = ETHER_HDR_LEN + sizeof(struct iphdr) + sizeof(struct tcphdr) + sizeof(struct ch_pkt_header);

    // 判断是否是 TLS 的 Client Hello 报文，
    // 这里规定了只抓 TLS 主版本为 1.0 的数据包，ssl v3 版本的报文不处理
    if( (ch->hsver != 771) || (ch->ct != 22) || (ch->hst != 1) ){
        /*
            todo
        */
        return;
    }

    inet_ntop(AF_INET, &(ip->saddr), client_ip, sizeof(client_ip));
    inet_ntop(AF_INET, &(ip->daddr), server_ip, sizeof(server_ip));

    // 获取 ja3 的TLS版本部分, 这里的TLS版本，不是TLS的主版本，是握手协议的版本
    snprintf(ver_str, VERSION_LEN, "%d", ch->hsver);

    // 新建一个字节流，读报文的 TLS 部分中的不具备固定格式的后半部分
    ByteReadStream tls_stream(p.get()->raw_data + pkt_read, (p.get()->pkt_info).caplen - pkt_read);

    // 跳过 Session ID 的值
    tls_stream.skip_len(ch->sidl);

    // 一定存在加密套件字段；如果不存在加密套件字段，说明这个报文存在问题
    // 不处理这样的 TLS Client Hello 报文
    if( !(tls_stream.check_len(2)) ){
        /*
            todo
        */
        return;
    }

    // 加密套件的个数
    uint32_t cipher_size = tls_stream.read_int16();
    cipher_size = cipher_size / 2;

    for(uint32_t i = 0; i < cipher_size; ++i)
    {
        uint32_t cipher = tls_stream.read_int16();
        if(check_grease(cipher))
            continue;
        
        if( ciphers_str_len )
            ciphers_str_len += snprintf(ciphers_str + ciphers_str_len, CIPHERS_LEN - ciphers_str_len, "%s", "-");
        
        ciphers_str_len += snprintf(ciphers_str + ciphers_str_len, CIPHERS_LEN - ciphers_str_len, "%d", cipher);
    }

    // 一定存在压缩方法字段；如果不存在压缩方法字段，说明这个报文存在问题
    // 不处理这样的 TLS Client Hello 报文
    if( !(tls_stream.check_len(1)) ){
        return;
    }

    // 压缩字段方法的长度通常为1
    uint32_t compression_len = tls_stream.read_int8();
    // 跳过压缩方法字段
    tls_stream.skip_len(compression_len);

    // TLS Client Hello 报文中，最后的部分都是扩展字段
    uint32_t exts_len = tls_stream.read_int16();

    // 把 读扩展字段 交给另一个 字节流来做
    ByteReadStream exts_stream(tls_stream, exts_len);
    // 同时用来读 tls 协议的流要跳过
    tls_stream.skip_len(exts_len);

    // 只知道扩展字段的总长度，并不知道扩展的个数，因此需要循环读取
    while ( !exts_stream.is_end() )
    {
        uint32_t ext_type = exts_stream.read_int16();
        uint32_t ext_len = exts_stream.read_int16();
        
        // 检查这个扩展是否是GREASE，是的话直接跳过
        if(check_grease(ext_type)){
            continue;
        }

        switch( ext_type )
        {
            case 0: // 这是 SNI 扩展字段
                if(ext_len == 0){ // 如果 SNI 扩展存在，但没有内容
                    sni_str_len = snprintf( sni_str, SNI_LEN, "%s", EMPTY_STR );
                }else{  // SNI 字段不为空
                    ByteReadStream sni_stream(exts_stream, ext_len);
                    exts_stream.skip_len(ext_len);
                    // 根据协议，一个 Clietn Hello 报文可能不止携带一个 “信息”, 
                    // 也可能携带除了域名之外的信息, 这里只要第一个信息，且只要域名信息
                    uint32_t sn_list_len = sni_stream.read_int16();
                    uint32_t sn_type = sni_stream.read_int8();
                    uint32_t sn_len = sni_stream.read_int16();
                    if (sn_type == 0)
                        sni_str_len = sni_stream.read_str(sni_str, sn_len, SNI_LEN);
                    else
                        sni_str_len = snprintf( sni_str, SNI_LEN, "%s", EMPTY_STR);
                }
                break;

            case 35:  // 这是 ticket 扩展字段
                if( ext_len == 0){  // 如果 Ticket 扩展存在，但没有内容
                    ticket_str_len = snprintf( ticket_str, TICKET_LEN, "%s", EMPTY_STR );
                }else{
                    ticket_str_len = exts_stream.read_str(ticket_str, ext_len, TICKET_LEN);
                }
                break;

            case 10:  // 这是 supported group 扩展字段
                if( ext_len == 0 ){     // 如果 supported group 扩展存在，但没有内容
                    supgs_str[0] = '\0';
                }else{   
                    ByteReadStream supgs_stream(exts_stream, ext_len);
                    exts_stream.skip_len(ext_len);
                    uint32_t supg_cnt = supgs_stream.read_int16() / 2;
                    for( uint32_t i = 0; i < supg_cnt; ++i ){
                        uint32_t supg = supgs_stream.read_int16();
                        if(check_grease(supg))
                            continue;
                        
                        if( supgs_str_len )
                            supgs_str_len += snprintf(supgs_str + supgs_str_len, SUPPG_LEN - supgs_str_len, "%s", "-");

                        supgs_str_len += snprintf(supgs_str + supgs_str_len, SUPPG_LEN - supgs_str_len, "%d", supg);                    
                    }
                }
                break;

            case 11:   // 这是 ec_point_formats 扩展字段
                if( ext_len == 0 ){    // 如果 support point format 扩展存在，但没有内容
                    suppfs_str[0] = '\0';
                }else{
                    ByteReadStream suppfs_stream(exts_stream, ext_len);
                    exts_stream.skip_len(ext_len);
                    uint32_t suppf_cnt = suppfs_stream.read_int8();
                    for( uint32_t i =0; i < suppf_cnt; ++i ){
                        uint32_t suppf = suppfs_stream.read_int8();

                        if( suppfs_str_len )
                            suppfs_str_len += snprintf(suppfs_str + suppfs_str_len, SUPPFS_LEN - suppfs_str_len, "%s", "-");
                        
                        suppfs_str_len += snprintf(suppfs_str + suppfs_str_len, SUPPFS_LEN - suppfs_str_len, "%d", suppf);
                    }
                }
                break;

            default:   // 其他类型的字段
                exts_stream.skip_len(ext_len);
        }

        if( exts_str_len )
            exts_str_len += snprintf(exts_str + exts_str_len, EXTS_LEN - exts_str_len, "%s", "-");
        
        exts_str_len += snprintf(exts_str + exts_str_len, EXTS_LEN - exts_str_len, "%d", ext_type);
    }

    // 扩展字段可以不存在，就是说读到的扩展字段长度可能为0，那么 exts_str 、supgs_str、suppfs_str 、sni_str、ticket_str 五个字符串为空
    if( exts_len == 0 ){
        exts_str[0] = '\0';
        supgs_str[0] = '\0';
        suppfs_str[0] = '\0';
        sni_str[0] = '\0';
        ticket_str[0] = '\0';
    }

    // 把 ver_str, ciphers_str, exts_str, supgs_str, suppfs_str 组合成 ja3
    ja3_str_len += snprintf(ja3_str + ja3_str_len, JA3_LEN, "%s,%s,%s,%s,%s", ver_str, ciphers_str, exts_str, supgs_str, suppfs_str);

    // 根据 ja3 字符串生成MD5哈希后的 ja3_md5
    unsigned char temp_md5[20];
    MD5_CTX ctx;

    MD5_Init(&ctx);
    MD5_Update(&ctx, ja3_str, ja3_str_len);
    MD5_Final(temp_md5, &ctx);
    for(int i=0; i < 16; i++){
        sprintf(ja3_md5 + i * 2, "%02X", temp_md5[i]);
    }
    ja3_md5[32] = '\0';

    // 把 ticket_str 转为16进制
    for(uint32_t i = 0; i < ticket_str_len; i++){
        sprintf(ticket_0X_str + i * 2, "%02X", ticket_str[i]);
    }

    // 把时间结果 struct timeval 转化为字符串。 timeval.tv_sec 为长整型 long int
    unix_time_len += snprintf(unix_time + unix_time_len, SMALL_STR_LEN - unix_time_len, "%ld", (p.get()->pkt_info).ts.tv_sec);
    unix_time_len += snprintf(unix_time + unix_time_len, SMALL_STR_LEN - unix_time_len, "%s%ld", ".", (p.get()->pkt_info).ts.tv_sec);
    // (p.get()->pkt_info).ts.tv_sec

    // 把最后的结果打包，写入队列中
    ch_info c_i(unix_time, server_ip, client_ip, ja3_md5, ticket_0X_str, sni_str);
    queue->push(c_i);
}